module GOLD_diag_to_Z

!***********************************************************************
!*                   GNU General Public License                        *
!* This file is a part of GOLD.                                        *
!*                                                                     *
!* GOLD is free software; you can redistribute it and/or modify it and *
!* are expected to follow the terms of the GNU General Public License  *
!* as published by the Free Software Foundation; either version 2 of   *
!* the License, or (at your option) any later version.                 *
!*                                                                     *
!* GOLD is distributed in the hope that it will be useful, but WITHOUT *
!* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY  *
!* or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public    *
!* License for more details.                                           *
!*                                                                     *
!* For the full text of the GNU General Public License,                *
!* write to: Free Software Foundation, Inc.,                           *
!*           675 Mass Ave, Cambridge, MA 02139, USA.                   *
!* or see:   http://www.gnu.org/licenses/gpl.html                      *
!***********************************************************************

!********+*********+*********+*********+*********+*********+*********+**
!*                                                                     *
!*  By Robert Hallberg, July 2006                                      *
!*                                                                     *
!*    This subroutine maps tracers and velocities into depth space     *
!*  for output as diagnostic quantities.  Currently, a piecewise       *
!*  linear subgrid structure is used for tracers, while velocities can *
!*  use either piecewise constant or piecewise linear structures.      *
!*                                                                     *
!*     A small fragment of the grid is shown below:                    *
!*                                                                     *
!*    j+1  x ^ x ^ x   At x:  q, f                                     *
!*    j+1  > o > o >   At ^:  v                                        *
!*    j    x ^ x ^ x   At >:  u                                        *
!*    j    > o > o >   At o:  h, D                                     *
!*    j-1  x ^ x ^ x                                                   *
!*        i-1  i  i+1  At x & ^:                                       *
!*           i  i+1    At > & o:                                       *
!*                                                                     *
!*  The boundaries always run through q grid points (x).               *
!*                                                                     *
!********+*********+*********+*********+*********+*********+*********+**

use GOLD_diag_mediator, only : post_data, register_diag_field, safe_alloc_ptr
use GOLD_diag_mediator, only : diag_ptrs, time_type, diag_axis_init
use GOLD_diag_mediator, only : ocean_register_diag
use GOLD_error_handler, only : GOLD_error, FATAL, WARNING
use GOLD_file_parser, only : get_param, log_param, log_version, param_file_type
use GOLD_grid, only : ocean_grid_type
use GOLD_io, only : slasher, vardesc
use GOLD_variables, only : p3d

use netcdf

implicit none ; private

#include <GOLD_memory.h>

public calculate_Z_diag_fields, register_Z_tracer, GOLD_diag_to_Z_init
public calculate_Z_transport, GOLD_diag_to_Z_end, ocean_register_diag_with_z
public find_overlap, find_limited_slope, register_Zint_diag, calc_Zint_diags

type, public :: diag_to_Z_CS ; private
  ! The following arrays are used to store diagnostics calculated in this
  ! module and unavailable outside of it.
  real, pointer, dimension(:,:,:) :: &
    u_z => NULL(), &    ! The zonal velocity remapped to depth space, in m s-1.
    v_z => NULL(), &    ! The meridional velocity remapped to depth space, in m s-1.
    uh_z => NULL(), &   ! The zonal transport remapped to depth space, in m3 s-1.
    vh_z => NULL()      ! The meridional transport remapped to depth space, in m3 s-1.
  type(p3d) :: tr_z(MAX_FIELDS)  ! An array of tracers, remapped to depth space.
  type(p3d) :: tr_model(MAX_FIELDS) ! Pointers to an array of tracers.
  real    :: missing_vel = -1.0e34
  real    :: missing_trans = -1.0e34
  real    :: missing_value = -1.0e34
  real    :: missing_tr(MAX_FIELDS) = -1.0e34
  integer :: id_u_z = -1, id_v_z = -1
  integer :: id_tr(MAX_FIELDS) = -1
  integer :: id_uh_Z = -1, id_vh_Z = -1
  integer :: num_tr_used = 0
  integer :: nz_zspace = -1
  real, pointer :: Z_int(:) => NULL()  ! The interface depths of the z-space file, in m.

  integer, dimension(3) :: axesqz, axeshz, axesuz, axesvz
  integer, dimension(3) :: axesqzi, axeshzi, axesuzi, axesvzi
  integer, dimension(1) :: axesz_out

  type(diag_ptrs), pointer :: diag ! A pointer to a structure of shareable
                             ! ocean diagnostic fields from other modules.

end type diag_to_Z_CS

integer, parameter :: NO_ZSPACE = -1

contains

subroutine calculate_Z_diag_fields(u, v, h, dt, G, CS)
  real, dimension(NXMEMQ_,NYMEM_,NZ_),  intent(in)    :: u
  real, dimension(NXMEM_,NYMEMQ_,NZ_),  intent(in)    :: v
  real, dimension(NXMEM_,NYMEM_,NZ_),   intent(in)    :: h
  real,                                 intent(in)    :: dt
  type(ocean_grid_type),                intent(inout) :: G
  type(diag_to_Z_CS),                   pointer       :: CS
!   This subroutine maps tracers and velocities into depth space for output.

! Arguments: u - Zonal velocity, in m s-1.
!  (in)      v - Meridional velocity, in m s-1.
!  (in)      h - Layer thickness, in m.
!  (in)      dt - the time difference in s since the last call to
!                 this subroutine.
!  (in)      G - The ocean's grid structure.
!  (in)      CS - The control structure returned by a previous call to
!                 diagnostics_init.
  real :: e(SZK_(G)+2)        ! The z-star interface heights in m.
  real :: h_f(SZK_(G)+1,SZI_(G)) ! The thicknesses of massive layers in m.
  real :: u_f(SZK_(G)+1,SZIQ_(G))! The zonal and meridional velocities in any
  real :: v_f(SZK_(G)+1,SZI_(G)) ! massive layers.
  real :: tr_f(SZK_(G),max(CS%num_tr_used,1),SZI_(G)) ! The tracer
                                 ! concentrations in massive layers.
    ! Note the deliberately reversed axes in h_f, u_f, v_f, and tr_f.
  integer :: nz_valid(SZIQ_(G))  ! The number of massive layers in a column.

  real :: D_pt(SZIQ_(G))       ! The bottom depth in m.
  real :: htot                 ! The summed layer thicknesses in m.
  real :: dilate               ! The proportion by which to dilate every layer.
  real :: wt(SZK_(G)+1)        ! The fractional weight for each layer in the
                               ! range between k_top and k_bot, nondim.
  real :: z1(SZK_(G)+1)        ! z1 and z2 are the depths of the top and bottom
  real :: z2(SZK_(G)+1)        ! limits of the part of a layer that contributes
                               ! to a depth level, relative to the cell center
                               ! and normalized by the cell thickness, nondim.
                               ! Note that -1/2 <= z1 < z2 <= 1/2.
  real :: sl_tr(max(CS%num_tr_used,1)) ! The normalized slope of the tracer
                                       ! within the cell, in tracer units.
  real :: slope ! A normalized slope of a variable within the cell.
  logical :: linear_velocity_profiles = .true.
  integer :: k_top, k_bot, k_bot_prev
  integer :: i, j, k, k2, kz, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, m, nkml
  is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = G%ke
  Isq = G%Iscq ; Ieq = G%Iecq ; Jsq = G%Jscq ; Jeq = G%Jecq
  nkml = max(G%nkml, 1)

  if (.not.associated(CS)) call GOLD_error(FATAL, &
         "diagnostic_fields_zstar: Module must be initialized before it is used.")

  ! If no fields are needed, return
  if ((CS%id_u_z <= 0) .and. (CS%id_v_z <= 0) .and. (CS%num_tr_used < 1)) return

  if (CS%id_u_z > 0) then
    do kz=1,CS%nz_zspace ; do j=js,je ; do I=Isq,Ieq
      CS%u_z(I,j,kz) = CS%missing_vel
    enddo ; enddo ; enddo

    do j=js,je
      ! Remove all massless layers.
      do I=Isq,Ieq ; nz_valid(I) = 0 ; D_pt(I) = 0.5*(G%D(i+1,j)+G%D(i,j)) ; enddo
      do k=1,nz ; do I=Isq,Ieq
        if ((G%umask(I,j) > 0.5) .and. (h(i,j,k)+h(i+1,j,k) > 4.0*G%Angstrom)) then
          nz_valid(I) = nz_valid(I) + 1 ; k2 = nz_valid(I)
          h_f(k2,I) = 0.5*(h(i,j,k)+h(i+1,j,k)) ; u_f(k2,I) = u(I,j,k)
        endif
      enddo ; enddo
      do I=Isq,Ieq ; if (G%umask(I,j) > 0.5) then
        ! Add an Angstrom thick layer at the bottom with 0 velocity to impose a
        ! no-slip BBC in the output, if anything but piecewise constant is used.
        nz_valid(I) = nz_valid(I) + 1 ; k2 = nz_valid(I)
        h_f(k2,I) = G%Angstrom ; u_f(k2,I) = 0.0
      endif ; enddo

      do I=Isq,Ieq ; if (nz_valid(I) > 0) then
      ! Calculate the z* interface heights for tracers.
        htot = 0.0 ; do k=1,nz_valid(i) ; htot = htot + h_f(k,i) ; enddo
        dilate = 0.0 ; if (htot*G%H_to_m > 0.5) dilate = (D_pt(i) - 0.0) / htot

        e(nz_valid(i)+1) = -D_pt(i)
        do k=nz_valid(i),1,-1 ; e(K) = e(K+1) + h_f(k,i)*dilate ; enddo

      ! Interpolate each variable into depth space.
        k_bot = 1 ; k_bot_prev = -1
        do kz=1,CS%nz_zspace
          call find_overlap(e, CS%Z_int(kz), CS%Z_int(kz+1), nz_valid(I), &
                            k_bot, k_top, k_bot, wt, z1, z2)
          if (k_top>nz_valid(I)) exit

          if (linear_velocity_profiles) then
            k = k_top
            if (k /= k_bot_prev) then
              ! Calculate the intra-cell profile.
              slope = 0.0 ! ; curv = 0.0
              if ((k < nz_valid(I)) .and. (k > nkml)) call &
                find_limited_slope(u_f(:,I), e, slope, k)
            endif
            ! This is the piecewise linear form.
            CS%u_z(I,j,kz) = wt(k) * (u_f(k,I) + 0.5*slope*(z2(k) + z1(k)))
            ! For the piecewise parabolic form add the following...
            !     + C1_3*curv*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            do k=k_top+1,k_bot-1
              CS%u_z(I,j,kz) = CS%u_z(I,j,kz) + wt(k)*u_f(k,I)
            enddo
            if (k_bot > k_top) then ; k = k_bot
              ! Calculate the intra-cell profile.
              slope = 0.0 ! ; curv = 0.0
              if ((k < nz_valid(I)) .and. (k > nkml)) call &
                find_limited_slope(u_f(:,I), e, slope, k)
               ! This is the piecewise linear form.
              CS%u_z(I,j,kz) = CS%u_z(I,j,kz) + wt(k) * &
                  (u_f(k,I) + 0.5*slope*(z2(k) + z1(k)))
              ! For the piecewise parabolic form add the following...
              !     + C1_3*curv*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            endif
            k_bot_prev = k_bot
          else ! Use piecewise constant profiles.
            CS%u_z(I,j,kz) = wt(k_top)*u_f(k_top,I)
            do k=k_top+1,k_bot
              CS%u_z(I,j,kz) = CS%u_z(I,j,kz) + wt(k)*u_f(k,I)
            enddo
          endif ! linear profiles
        enddo ! kz-loop
      endif ; enddo ! I-loop and mask
    enddo ! j-loop
    
    call post_data(CS%id_u_z, CS%u_z, CS%diag)
  endif

  if (CS%id_v_z > 0) then
    do kz=1,CS%nz_zspace ; do J=Jsq,Jeq ; do i=is,ie
      CS%v_z(i,J,kz) = CS%missing_vel
    enddo ; enddo ; enddo

    do J=Jsq,Jeq
      ! Remove all massless layers.
      do i=is,ie ; nz_valid(i) = 0 ; D_pt(i) = 0.5*(G%D(i,j)+G%D(i,j+1)) ; enddo
      do k=1,nz ; do i=is,ie
        if ((G%vmask(i,j) > 0.5) .and. (h(i,j,k)+h(i,j+1,k) > 4.0*G%Angstrom)) then
          nz_valid(i) = nz_valid(i) + 1 ; k2 = nz_valid(i)
          h_f(k2,i) = 0.5*(h(i,j,k)+h(i,j+1,k)) ; v_f(k2,i) = v(i,j,k)
        endif
      enddo ; enddo
      do i=is,ie ; if (G%vmask(i,j) > 0.5) then
        ! Add an Angstrom thick layer at the bottom with 0 velocity to impose a
        ! no-slip BBC in the output, if anything but piecewise constant is used.
        nz_valid(i) = nz_valid(i) + 1 ; k2 = nz_valid(i)
        h_f(k2,i) = G%Angstrom ; v_f(k2,i) = 0.0
      endif ; enddo

      do i=is,ie ; if (nz_valid(i) > 0) then
      ! Calculate the z* interface heights for tracers.
        htot = 0.0 ; do k=1,nz_valid(i) ; htot = htot + h_f(k,i) ; enddo
        dilate = 0.0 ; if (htot > 0.5) dilate = (D_pt(i) - 0.0) / htot

        e(nz_valid(i)+1) = -D_pt(i)
        do k=nz_valid(i),1,-1 ; e(K) = e(K+1) + h_f(k,i)*dilate ; enddo

      ! Interpolate each variable into depth space.
        k_bot = 1 ; k_bot_prev = -1
        do kz=1,CS%nz_zspace
          call find_overlap(e, CS%Z_int(kz), CS%Z_int(kz+1), nz_valid(i), &
                            k_bot, k_top, k_bot, wt, z1, z2)
          if (k_top>nz_valid(i)) exit

          if (linear_velocity_profiles) then
            k = k_top
            if (k /= k_bot_prev) then
              ! Calculate the intra-cell profile.
              slope = 0.0 ! ; curv = 0.0
              if ((k < nz_valid(i)) .and. (k > nkml)) call &
                find_limited_slope(v_f(:,i), e, slope, k)
            endif
            ! This is the piecewise linear form.
            CS%v_z(i,J,kz) = wt(k) * (v_f(k,i) + 0.5*slope*(z2(k) + z1(k)))
            ! For the piecewise parabolic form add the following...
            !     + C1_3*curv*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            do k=k_top+1,k_bot-1
              CS%v_z(i,J,kz) = CS%v_z(i,J,kz) + wt(k)*v_f(k,i)
            enddo
            if (k_bot > k_top) then ; k = k_bot
              ! Calculate the intra-cell profile.
              slope = 0.0 ! ; curv = 0.0
              if ((k < nz_valid(i)) .and. (k > nkml)) call &
                find_limited_slope(v_f(:,i), e, slope, k)
               ! This is the piecewise linear form.
              CS%v_z(i,J,kz) = CS%v_z(i,J,kz) + wt(k) * &
                  (v_f(k,i) + 0.5*slope*(z2(k) + z1(k)))
              ! For the piecewise parabolic form add the following...
              !     + C1_3*curv*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            endif
            k_bot_prev = k_bot
          else ! Use piecewise constant profiles.
            CS%v_z(i,J,kz) = wt(k_top)*v_f(k_top,i)
            do k=k_top+1,k_bot
              CS%v_z(i,J,kz) = CS%v_z(i,J,kz) + wt(k)*v_f(k,i)
            enddo
          endif ! linear profiles
        enddo ! kz-loop
      endif ; enddo ! i-loop and mask
    enddo ! J-loop
    
    call post_data(CS%id_v_z, CS%v_z, CS%diag)
  endif

  if (CS%num_tr_used > 0) then
  
    do m=1,CS%num_tr_used ; do kz=1,CS%nz_zspace ; do j=js,je ; do i=is,ie
      CS%tr_z(m)%p(i,j,kz) = CS%missing_tr(m)
    enddo ; enddo ; enddo ; enddo

    do j=js,je
      ! Remove all massless layers.
      do i=is,ie ; nz_valid(i) = 0 ; D_pt(i) = G%D(i,j) ; enddo
      do k=1,nz ; do i=is,ie
        if ((G%hmask(i,j) > 0.5) .and. (h(i,j,k) > 2.0*G%Angstrom)) then
          nz_valid(i) = nz_valid(i) + 1 ; k2 = nz_valid(i)
          h_f(k2,i) = h(i,j,k)
          do m=1,CS%num_tr_used ; tr_f(k2,m,i) = CS%tr_model(m)%p(i,j,k) ; enddo
        endif
      enddo ; enddo

      do i=is,ie ; if (nz_valid(i) > 0) then
      ! Calculate the z* interface heights for tracers.
        htot = 0.0 ;  do k=1,nz_valid(i) ; htot = htot + h_f(k,i) ; enddo
        dilate = 0.0 ; if (htot > 0.5) dilate = (D_pt(i) - 0.0) / htot

        e(nz_valid(i)+1) = -D_pt(i)
        do k=nz_valid(i),1,-1 ; e(K) = e(K+1) + h_f(k,i)*dilate ; enddo

      ! Interpolate each variable into depth space.
        k_bot = 1 ; k_bot_prev = -1
        do kz=1,CS%nz_zspace
          call find_overlap(e, CS%Z_int(kz), CS%Z_int(kz+1), nz_valid(i), &
                            k_bot, k_top, k_bot, wt, z1, z2)
          if (k_top>nz_valid(i)) exit

          do m=1,CS%num_tr_used
            k = k_top
            if (k /= k_bot_prev) then
              ! Calculate the intra-cell profile.
              sl_tr(m) = 0.0 ! ; cur_tr(m) = 0.0
              if ((k < nz_valid(i)) .and. (k > nkml)) call &
                find_limited_slope(tr_f(:,m,i), e, sl_tr(m), k)
            endif
            ! This is the piecewise linear form.
            CS%tr_z(m)%p(i,j,kz) = wt(k) * &
                (tr_f(k,m,i) + 0.5*sl_tr(m)*(z2(k) + z1(k)))
            ! For the piecewise parabolic form add the following...
            !     + C1_3*cur_tr(m)*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            do k=k_top+1,k_bot-1
              CS%tr_z(m)%p(i,j,kz) = CS%tr_z(m)%p(i,j,kz) + wt(k)*tr_f(k,m,i)
            enddo
            if (k_bot > k_top) then
              k = k_bot
              ! Calculate the intra-cell profile.
              sl_tr(m) = 0.0 ! ; cur_tr(m) = 0.0
              if ((k < nz_valid(i)) .and. (k > nkml)) call &
                find_limited_slope(tr_f(:,m,i), e, sl_tr(m), k)
              ! This is the piecewise linear form.
              CS%tr_z(m)%p(i,j,kz) = CS%tr_z(m)%p(i,j,kz) + wt(k) * &
                  (tr_f(k,m,i) + 0.5*sl_tr(m)*(z2(k) + z1(k)))
              ! For the piecewise parabolic form add the following...
              !     + C1_3*cur_tr(m)*(z2(k)**2 + z2(k)*z1(k) + z1(k)**2))
            endif
          enddo
          k_bot_prev = k_bot
        enddo ! kz-loop
      endif ; enddo ! i-loop and mask

    enddo ! j-loop
    
    do m=1,CS%num_tr_used
      if (CS%id_tr(m) > 0) call post_data(CS%id_tr(m), CS%tr_z(m)%p, CS%diag)
    enddo
  endif

end subroutine calculate_Z_diag_fields

subroutine calculate_Z_transport(uh_int, vh_int, h, dt, G, CS)
  real, dimension(NXMEMQ_,NYMEM_,NZ_),  intent(in)    :: uh_int
  real, dimension(NXMEM_,NYMEMQ_,NZ_),  intent(in)    :: vh_int
  real, dimension(NXMEM_,NYMEM_,NZ_),   intent(in)    :: h
  real,                                 intent(in)    :: dt
  type(ocean_grid_type),                intent(inout) :: G
  type(diag_to_Z_CS),                   pointer       :: CS
!   This subroutine maps tracers and velocities into depth space for output.

! Arguments: uh_int - Time integrated zonal transport, in m3.
!  (in)      vh_int - Time integrated meridional transport, in m3.
!  (in)      h - Layer thickness, in m.
!  (in)      dt - the time difference in s since the last call to
!                 this subroutine.
!  (in)      G - The ocean's grid structure.
!  (in)      CS - The control structure returned by a previous call to
!                 diagnostics_init.
  real, dimension(SZI_(G), SZJ_(G)) :: &
    htot, &        ! The total layer thickness in m.
    dilate         ! A nondimensional factor by which to dilate layers to
                   ! convert them into z* space.  (-G%D < z* < 0)
  real, dimension(SZI_(G), max(CS%nz_zspace,1)) :: &
    uh_Z           ! uh_int interpolated into depth space, in m3.
  real, dimension(SZIQ_(G), max(CS%nz_zspace,1)) :: &
    vh_Z           ! vh_int interpolated into depth space, in m3.
  real :: h_rem    ! The dilated thickness of a layer that has yet to be mapped
                   ! into depth space, in m.
  real :: uh_rem   ! The integrated zonal transport of a layer that has yet to be
                   ! mapped into depth space, in m3.
  real :: vh_rem   ! The integrated meridional transport of a layer that has yet
                   ! to be mapped into depth space, in m3.
  real :: h_here   ! The thickness of a layer that is within the range of the
                   ! current depth level, in m.
  real :: h_above  ! The thickness of a layer that is above the current depth
                   ! level, in m.
  real :: uh_here  ! The zonal transport of a layer that is attributed to the
                   ! current depth level, in m3.
  real :: vh_here  ! The meridional transport of a layer that is attributed to
                   ! the current depth level, in m3.
  real :: Idt      ! The inverse of the time step in s.
  real :: Z_int_above(SZIQ_(G)) ! The height of the interface atop a layer, m.
  integer :: kz(SZIQ_(G)) ! The index of the depth level that is being
                 ! contributed to.
  integer :: i, j, k, is, ie, js, je, Isq, Ieq, Jsq, Jeq, nz, nz_z
  is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = G%ke
  Isq = G%Iscq ; Ieq = G%Iecq ; Jsq = G%Jscq ; Jeq = G%Jecq

  if (.not.associated(CS)) call GOLD_error(FATAL, &
         "calculate_Z_transport: Module must be initialized before it is used.")
  if ((CS%id_uh_Z <= 0) .and. (CS%id_vh_Z <= 0)) return

  Idt = 1.0 ; if (dt > 0.0) Idt = 1.0 / dt
  nz_z = CS%nz_zspace

  if (nz_z <= 0) return

  ! Determine how much the layers will be dilated in recasting them into z*
  ! coordiantes.  (-G%D < z* < 0).
  do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1
    htot(i,j) = G%H_subroundoff
  enddo ; enddo
  do k=1,nz ; do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1
    htot(i,j) = htot(i,j) + h(i,j,k)
  enddo ; enddo ; enddo
  do j=Jsq,Jeq+1 ; do i=Isq,Ieq+1
    dilate(i,j) = G%D(i,j) / htot(i,j)
  enddo ; enddo

  if (CS%id_uh_Z > 0) then ; do j=js,je
    do I=Isq,Ieq
      kz(I) = nz_z ; z_int_above(I) = -0.5*(G%D(i,j)+G%D(i+1,j))
    enddo
    do k=nz_z,1,-1 ; do I=Isq,Ieq
      uh_Z(I,k) = 0.0
      if (CS%Z_int(k) < z_int_above(I)) kz(I) = k-1
    enddo ; enddo
    do k=nz,1,-1 ; do I=Isq,Ieq
      h_rem = 0.5*(dilate(i,j)*h(i,j,k) + dilate(i+1,j)*h(i+1,j,k))
      uh_rem = uh_int(I,j,k)
      z_int_above(I) = z_int_above(I) + h_rem

      do ! Distribute this layer's transport into the depth-levels.
        h_above = z_int_above(I) - CS%Z_int(kz(I)) 
        if ((kz(I) == 1) .or. (h_above <= 0.0) .or. (h_rem <= 0.0)) then
          ! The entire remaining transport is on this level.
          uh_Z(I,kz(I)) = uh_Z(I,kz(I)) + uh_rem ; exit
        else
          h_here = h_rem - h_above
          uh_here = uh_rem * (h_here / h_rem)

          h_rem = h_rem - h_here ! = h_above
          uh_Z(I,kz(I)) = uh_Z(I,kz(I)) + uh_here
          uh_rem = uh_rem - uh_here
          kz(I) = kz(I) - 1
        endif
      enddo ! End of loop through the target depth-space levels.
    enddo ; enddo
    do k=1,nz_z ; do I=Isq,Ieq
      CS%uh_z(I,j,k) = uh_Z(I,k)*Idt
    enddo ; enddo
  enddo ; endif
  if (CS%id_vh_Z > 0) then ; do J=Jsq,Jeq
    do i=is,ie
      kz(i) = nz_z ; z_int_above(i) = -0.5*(G%D(i,j)+G%D(i,j+1))
    enddo
    do k=nz_z,1,-1 ; do i=is,ie
      vh_Z(i,k) = 0.0
      if (CS%Z_int(k) < z_int_above(i)) kz(i) = k-1
    enddo ; enddo
    do k=nz,1,-1 ; do i=is,ie
      h_rem = 0.5*(dilate(i,j)*h(i,j,k) + dilate(i,j+1)*h(i,j+1,k))
      vh_rem = vh_int(i,J,k)
      z_int_above(i) = z_int_above(i) + h_rem

      do ! Distribute this layer's transport into the depth-levels.
        h_above = z_int_above(i) - CS%Z_int(kz(i)) 
        if ((kz(i) == 1) .or. (h_above <= 0.0) .or. (h_rem <= 0.0)) then
          ! The entire remaining transport is on this level.
          vh_Z(i,kz(i)) = vh_Z(i,kz(i)) + vh_rem ; exit
        else
          h_here = h_rem - h_above
          vh_here = vh_rem * (h_here / h_rem)

          h_rem = h_rem - h_here ! = h_above
          vh_Z(i,kz(i)) = vh_Z(i,kz(i)) + vh_here
          vh_rem = vh_rem - vh_here
          kz(i) = kz(i) - 1
        endif
      enddo ! End of loop through the target depth-space levels.
    enddo ; enddo
    do k=1,nz_z ; do i=is,ie
      CS%vh_z(i,J,k) = vh_Z(i,k)*Idt
    enddo ; enddo
  enddo ; endif

  if (CS%id_uh_Z > 0) call post_data(CS%id_uh_Z, CS%uh_z, CS%diag)
  if (CS%id_vh_Z > 0) call post_data(CS%id_vh_Z, CS%vh_z, CS%diag)

end subroutine calculate_Z_transport

subroutine find_overlap(e, Z_top, Z_bot, k_max, k_start, k_top, k_bot, wt, z1, z2)
  real, dimension(:), intent(in) :: e
  real, intent(in)   :: Z_top, Z_bot
  integer, intent(in) :: k_max, k_start
  integer, intent(out) :: k_top, k_bot
  real, dimension(:), intent(out) :: wt, z1, z2

!   This subroutine determines the layers bounded by interfaces e that overlap
! with the depth range between Z_top and Z_bot, and also the fractional weights
! of each layer. It also calculates the normalized relative depths of the range
! of each layer that overlaps that depth range.
!   Note that by convention, e decreases with increasing k and Z_top > Z_bot.
!
! Arguments: e - A column's interface heights, in m.
!  (in)      Z_top - The top of the range being mapped to, in m.
!  (in)      Z_bot - The bottom of the range being mapped to, in m.
!  (in)      k_max - The number of valid layers.
!  (in)      k_start - The layer at which to start searching.
!  (out)     k_top, k_bot - The indices of the top and bottom layers that
!                           overlap with the depth range.
!  (out)     wt - The relative weights of each layer from k_top to k_bot.
!  (out)     z1, z2 - z1 and z2 are the depths of the top and bottom limits of
!                     the part of a layer that contributes to a depth level, 
!                     relative to the cell center and normalized by the cell
!                     thickness, nondim.  Note that -1/2 <= z1 < z2 <= 1/2.
  real :: Ih, e_c, tot_wt, I_totwt
  integer :: k
  
  do k=k_start,k_max ; if (e(K+1)<Z_top) exit ; enddo
  k_top = k
  if (k>k_max) return

  ! Determine the fractional weights of each layer.
  ! Note that by convention, e and Z_int decrease with increasing k.
  if (e(K+1)<=Z_bot) then
    wt(k) = 1.0 ; k_bot = k
    Ih = 1.0 / (e(K)-e(K+1))
    e_c = 0.5*(e(K)+e(K+1))
    z1(k) = (e_c - MIN(e(K),Z_top)) * Ih
    z2(k) = (e_c - Z_bot) * Ih
  else
    wt(k) = MIN(e(K),Z_top) - e(K+1) ; tot_wt = wt(k) ! These are always > 0.
    z1(k) = (0.5*(e(K)+e(K+1)) - MIN(e(K),Z_top)) / (e(K)-e(K+1))
    z2(k) = 0.5
    k_bot = k_max
    do k=k_top+1,k_max
      if (e(K+1)<=Z_bot) then
        k_bot = k
        wt(k) = e(K) - Z_bot ; z1(k) = -0.5
        z2(k) = (0.5*(e(K)+e(K+1)) - Z_bot) / (e(K)-e(K+1))
      else
        wt(k) = e(K) - e(K+1) ; z1(k) = -0.5 ; z2(k) = 0.5
      endif
      tot_wt = tot_wt + wt(k) ! wt(k) is always > 0.
      if (k>=k_bot) exit
    enddo

    I_totwt = 1.0 / tot_wt
    do k=k_top,k_bot ; wt(k) = I_totwt*wt(k) ; enddo
  endif
  
end subroutine find_overlap
  

subroutine find_limited_slope(val, e, slope, k)
  real, dimension(:), intent(in) :: val
  real, dimension(:), intent(in) :: e
  real, intent(out) :: slope
  integer, intent(in) :: k
!   This subroutine determines a limited slope for val to be advected with
! a piecewise limited scheme.

! Arguments: val - An column the values that are being interpolated.
!  (in)      e - A column's interface heights, in m.
!  (in)      slope - The normalized slope in the intracell distribution of val.
!  (in)      k - The layer whose slope is being determined.
  real :: d1, d2

  if ((val(k)-val(k-1)) * (val(k)-val(k+1)) >= 0.0) then
    slope = 0.0 ! ; curvature = 0.0   
  else
    d1 = 0.5*(e(K-1)-e(K+1)) ; d2 = 0.5*(e(K)-e(K+2))
    slope = (d1**2*(val(k+1) - val(k)) + d2**2*(val(k) - val(k-1))) * &
            ((e(K) - e(K+1)) / (d1*d2*(d1+d2)))
    ! slope = 0.5*(val(k+1) - val(k-1))
    ! This is S.J. Lin's form of the PLM limiter.
    slope = sign(1.0,slope) * min(abs(slope), &
        2.0*(max(val(k-1),val(k),val(k+1)) - val(k)), &
        2.0*(val(k) - min(val(k-1),val(k),val(k+1))))
    ! curvature = 0.0  
  endif

end subroutine find_limited_slope

subroutine calc_Zint_diags(h, in_ptrs, ids, num_diags, &
                                       G, CS)
  real, dimension(NXMEM_,NYMEM_,NZ_), intent(in) :: h
  type(p3d), dimension(:),            intent(in) :: in_ptrs
  integer,   dimension(:),            intent(in) :: ids
  integer,                            intent(in) :: num_diags
  type(ocean_grid_type),              intent(in) :: G
  type(diag_to_Z_CS),                 pointer    :: CS
  
  real, dimension(SZI_(G),SZJ_(G),max(CS%nz_zspace+1,1),max(num_diags,1)) :: &
    diag_on_Z  ! The diagnostics interpolated to depth space.
  real, dimension(SZI_(G),SZK_(G)+1) :: e
  real, dimension(max(num_diags,1),SZI_(G),SZK_(G)+1) :: diag2d
  real, dimension(SZI_(G)) :: &
    htot, &              ! The summed layer thicknesses in m or kg m-2.
    dilate               ! The proportion by which to dilate every layer.
  real :: wt             ! The weighting of the interface above in the
                         ! interpolation to the target depths.
  integer :: kL(SZI_(G)) ! The layer-space index of the shallowest interface
                         ! below the target depth.
  integer :: i, j, k, k2, kz, is, ie, js, je, nz, m
  is = G%isc ; ie = G%iec ; js = G%jsc ; je = G%jec ; nz = G%ke
    
  if (num_diags < 1) return
  if (.not.associated(CS)) call GOLD_error(FATAL, &
       "calc_Zint_diags: Module must be initialized before it is used.")
  
  do j=js,je
    ! Calculate the stretched z* interface depths.
    do i=is,ie ; htot(i) = 0.0 ; kL(i) = 1 ; enddo
    do k=1,nz ; do i=is,ie ; htot(i) = htot(i) + h(i,j,k) ; enddo ; enddo
    do i=is,ie
      dilate(i) = 0.0
      if (htot(i)*G%H_to_m > 0.5) dilate(i) = (G%D(i,j) - 0.0) / htot(i)
      e(i,nz+1) = -G%D(i,j)
    enddo
    do k=nz,1,-1 ; do i=is,ie
      e(i,k) = e(i,k+1) + h(i,j,k) * dilate(i)
    enddo ; enddo
    ! e(i,1) should be 0 as a consistency check.

    do k=1,nz+1 ; do i=is,ie ; do m=1,num_diags
      diag2d(m,i,k) = in_ptrs(m)%p(i,j,k)
    enddo ; enddo ; enddo

    do kz=1,CS%nz_zspace+1 ; do i=is,ie
      ! Find the interface below the target Z-file depth, kL.
      if (CS%Z_int(kz) < e(i,nz+1)) then
        kL(i) = nz+2
      else
        do k=kL(i),nz+1 ; if (CS%Z_int(kz) > e(i,k)) exit ; enddo
        kL(i) = k
      endif
      if (kL(i)>1) then
        if (CS%Z_int(kz) > e(i,kL(i)-1)) call GOLD_error(FATAL, &
        "calc_Zint_diags: Interface depth mapping is incorrect.")
      endif
      if ((kL(i)>1) .and. (kL(i)<=nz+1)) then
        if (e(i,kL(i)-1) == e(i,kL(i))) call GOLD_error(WARNING, &
          "calc_Zint_diags: Interface depths equal.", all_print=.true.)
        if (e(i,kL(i)-1) - e(i,kL(i)) < 0.0) call GOLD_error(FATAL, &
          "calc_Zint_diags: Interface depths inverted.")
      endif

      if (kL(i) <= 1) then
        do m=1,num_diags
          diag_on_Z(i,j,kz,m) = diag2d(m,i,1)
        enddo
      elseif (kL(i) > nz+1) then
        do m=1,num_diags
          diag_on_Z(i,j,kz,m) = CS%missing_value
        enddo
      else
        wt = 0.0 ! This probably should not happen?
        if (e(i,kL(i)-1) - e(i,kL(i)) > 0.0) &
          wt = (CS%Z_int(kz) - e(i,kL(i))) / (e(i,kL(i)-1) - e(i,kL(i)))
        if ((wt < 0.0) .or. (wt > 1.0)) call GOLD_error(FATAL, &
          "calc_Zint_diags: Bad value of wt found.")
        do m=1,num_diags
          diag_on_Z(i,j,kz,m) = wt * diag2d(m,i,kL(i)-1) + &
                                (1.0-wt) * diag2d(m,i,kL(i))
        enddo
      endif
    enddo ; enddo

  enddo

  do m=1,num_diags
    if (ids(m) > 0) call post_data(ids(m), diag_on_Z(:,:,:,m), CS%diag)
  enddo

end subroutine calc_Zint_diags

                          
subroutine register_Z_tracer(tr_ptr, name, long_name, units, Time, G, CS)
  real, dimension(NXMEM_,NYMEM_,NZ_), target, intent(in) :: tr_ptr
  character(len=*),                           intent(in) :: name
  character(len=*),                           intent(in) :: long_name
  character(len=*),                           intent(in) :: units
  type(time_type),                            intent(in) :: Time
  type(ocean_grid_type),                      intent(in) :: G
  type(diag_to_Z_CS),                         pointer    :: CS
!   This subroutine registers a tracer to be output in depth space.
! Arguments: tr_ptr - the tracer that is being offered for translation to Z-space.
!  (in)      name - The name to be used for the output tracer.
!  (in)      long_name - The long name to be used for the output tracer.
!  (in)      units - The units of the output tracer.
!  (in)      Time - The current model time.
!  (in)      G - The ocean's grid structure.
!  (in)      CS - The control structure returned by a previous call to
!                 diagnostics_init.
  integer :: isd, ied, jsd, jed, nz, m, id_test
  isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed ; nz = G%ke
  if (.not.associated(CS)) call GOLD_error(FATAL, &
         "register_Z_tracer: Module must be initialized before it is used.")

  if (CS%num_tr_used >= MAX_FIELDS) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z:  Attempted to register and use "//&
                   "more than MAX_FIELDS z-space tracers via register_Z_tracer.")
    return
  endif

  m = CS%num_tr_used + 1

  CS%missing_tr(m) = CS%missing_value ! This could be changed later, if desired.
  if (CS%nz_zspace > 0) then
    CS%id_tr(m) = register_diag_field('ocean_model', name, CS%axeshz, Time, &
                                      long_name, units, missing_value=CS%missing_tr(m))
  else
    id_test = register_diag_field('ocean_model', name, G%axesh1, Time, &
                                  long_name, units, missing_value=CS%missing_tr(m))
    if (id_test>0) call GOLD_error(WARNING, &
        "GOLD_diag_to_Z_init: "//trim(name)// &
        " cannot be output without an appropriate depth-space target file.")
  endif

  if (CS%id_tr(m) <= 0) then ; CS%id_tr(m) = -1 ; return ; endif

  CS%num_tr_used = m  
  call safe_alloc_ptr(CS%tr_z(m)%p,isd,ied,jsd,jed,CS%nz_zspace)
  CS%tr_model(m)%p => tr_ptr

end subroutine register_Z_tracer

subroutine GOLD_diag_to_Z_init(Time, G, param_file, diag, CS)
  type(time_type),         intent(in)    :: Time
  type(ocean_grid_type),   intent(in)    :: G
  type(param_file_type),   intent(in)    :: param_file
  type(diag_ptrs), target, intent(inout) :: diag
  type(diag_to_Z_CS),      pointer       :: CS
! Arguments: Time - The current model time.
!  (in)      G - The ocean's grid structure.
!  (in)      param_file - A structure indicating the open file to parse for
!                         model parameter values.
!  (in)      diag - A structure containing pointers to common diagnostic fields.
!  (in/out)  CS - A pointer that is set to point to the control structure
!                 for this module
  character(len=128) :: version = '$Id$'
  character(len=128) :: tagname = '$Name$'
  character(len=40)  :: mod = "GOLD_diag_to_Z" ! This module's name.
  character(len=200) :: in_dir, zgrid_file   ! Strings for directory/file.
  character(len=48)  :: flux_units
  integer :: z_axis, zint_axis
  integer :: isd, ied, jsd, jed, Isdq, Iedq, Jsdq, Jedq, nz, id_test
  isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed ; nz = G%ke
  Isdq = G%Isdq ; Iedq = G%Iedq ; Jsdq = G%Jsdq ; Jedq = G%Jedq

  if (associated(CS)) then
    call GOLD_error(WARNING, "GOLD_diag_to_Z_init called with an associated "// &
                            "control structure.")
    return
  endif
  allocate(CS)

  if (G%Boussinesq) then ; flux_units = "meter3 second-1"
  else ; flux_units = "kilogram second-1" ; endif

  CS%diag => diag

  ! Read all relevant parameters and write them to the model log.
  call log_version(param_file, mod, version, tagname)
  ! Read in z-space info from a NetCDF file.
  call get_param(param_file, mod, "Z_OUTPUT_GRID_FILE", zgrid_file, &
                 "The file that specifies the vertical grid for \n"//&
                 "depth-space diagnostics, or blank to disable \n"//&
                 "depth-space output.", default="")
  if (len_trim(zgrid_file) > 0) then
    call get_param(param_file, mod, "INPUTDIR", in_dir, &
                 "The directory in which input files are found.", default=".")
    in_dir = slasher(in_dir)  
    call get_Z_depths(trim(in_dir)//trim(zgrid_file), "zw", CS%Z_int, "zt", &
                      z_axis, zint_axis, CS%nz_zspace)
    call log_param(param_file, mod, "INPUTDIR/Z_OUTPUT_GRID_FILE", &
                   trim(in_dir)//trim(zgrid_file))
    call log_param(param_file, mod, "NZ_ZSPACE (from file)", CS%nz_zspace, &
                 "The number of depth-space levels.  This is determined \n"//&
                 "from the size of the variable zw in the output grid file.", &
                 units="nondim")
  else
    in_dir = "" ; CS%nz_zspace = -1
  endif

  if (CS%nz_zspace > 0) then
    CS%axesqz = (/ G%axesq1(1), G%axesq1(2), z_axis /)
    CS%axeshz = (/ G%axesh1(1), G%axesh1(2), z_axis /)
    CS%axesuz = (/ G%axesu1(1), G%axesu1(2), z_axis /)
    CS%axesvz = (/ G%axesv1(1), G%axesv1(2), z_axis /)
    CS%axesqzi = (/ G%axesq1(1), G%axesq1(2), zint_axis /)
    CS%axeshzi = (/ G%axesh1(1), G%axesh1(2), zint_axis /)
    CS%axesuzi = (/ G%axesu1(1), G%axesu1(2), zint_axis /)
    CS%axesvzi = (/ G%axesv1(1), G%axesv1(2), zint_axis /)

    CS%id_u_z = register_diag_field('ocean_model', 'u_z', CS%axesuz, Time, &
        'Zonal Velocity in Depth Space', 'meter second-1', &
        missing_value=CS%missing_vel)
    if (CS%id_u_z>0) call safe_alloc_ptr(CS%u_z,Isdq,Iedq,jsd,jed,CS%nz_zspace)

    CS%id_v_z = register_diag_field('ocean_model', 'v_z', CS%axesvz, Time, &
        'Meridional Velocity in Depth Space', 'meter second-1', &
        missing_value=CS%missing_vel)
    if (CS%id_v_z>0) call safe_alloc_ptr(CS%v_z,isd,ied,Jsdq,Jedq,CS%nz_zspace)

    CS%id_uh_z = register_diag_field('ocean_model', 'uh_z', CS%axesuz, Time, &
        'Zonal Volume Transport in Depth Space', flux_units, &
        missing_value=CS%missing_trans)
    if (CS%id_uh_z>0) call safe_alloc_ptr(CS%uh_z,Isdq,Iedq,jsd,jed,CS%nz_zspace)

    CS%id_vh_z = register_diag_field('ocean_model', 'vh_z', CS%axesvz, Time, &
        'Meridional Volume Transport in Depth Space', flux_units, &
        missing_value=CS%missing_trans)
    if (CS%id_vh_z>0) call safe_alloc_ptr(CS%vh_z,isd,ied,Jsdq,Jedq,CS%nz_zspace)

  else
    ! Check whether the diag-table is requesting any z-space files, and issue
    ! a warning if it is.
    id_test = register_diag_field('ocean_model', 'u_z', G%axesu1, Time, &
        'Zonal Velocity in Depth Space', 'meter second-1')
    if (id_test>0) call GOLD_error(WARNING, &
        "GOLD_diag_to_Z_init: u_z cannot be output without "//&
        "an appropriate depth-space target file.")

    id_test = register_diag_field('ocean_model', 'v_z', G%axesv1, Time, &
        'Meridional Velocity in Depth Space', 'meter second-1')
    if (id_test>0) call GOLD_error(WARNING, &
        "GOLD_diag_to_Z_init: v_z cannot be output without "//&
        "an appropriate depth-space target file.")

    id_test = register_diag_field('ocean_model', 'uh_z', G%axesu1, Time, &
        'Meridional Volume Transport in Depth Space', flux_units)
    if (id_test>0) call GOLD_error(WARNING, &
        "GOLD_diag_to_Z_init: uh_z cannot be output without "//&
        "an appropriate depth-space target file.")

    id_test = register_diag_field('ocean_model', 'vh_z', G%axesv1, Time, &
        'Meridional Volume Transport in Depth Space', flux_units)
    if (id_test>0) call GOLD_error(WARNING, &
        "GOLD_diag_to_Z_init: vh_z cannot be output without "//&
        "an appropriate depth-space target file.")
  endif

end subroutine GOLD_diag_to_Z_init

subroutine get_Z_depths(depth_file, int_depth_name, int_depth, cell_depth_name, &
                        z_axis_index, edge_index, nz_out)
  character(len=*), intent(in) :: depth_file
  character(len=*), intent(in) :: int_depth_name
  real, dimension(:), pointer  :: int_depth
  character(len=*), intent(in) :: cell_depth_name
  integer, intent(out) :: z_axis_index, edge_index
  integer, intent(out) :: nz_out
!   This subroutine reads the depths of the interfaces bounding the intended
! layers from a NetCDF file.  If no appropriate file is found, -1 is returned
! as the number of layers in the output file.  Also, a diag_manager axis is set
! up with the same information as this axis.
  real, allocatable :: cell_depth(:)
  character (len=200) :: units, long_name
  integer :: ncid, status, intid, intvid, layid, layvid, k, ni

  nz_out = -1
  
  status = NF90_OPEN(depth_file, NF90_NOWRITE, ncid);
  if (status /= NF90_NOERR) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
        " Difficulties opening "//trim(depth_file)//" - "//&
        trim(NF90_STRERROR(status)))
    nz_out = -1 ; return
  endif

  status = NF90_INQ_DIMID(ncid, int_depth_name, intid)
  if (status /= NF90_NOERR) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting ID of dimension "//&
      trim(int_depth_name)//" in "//trim(depth_file))
    nz_out = -1 ; return
  endif
  status = nf90_Inquire_Dimension(ncid, intid, len=ni)
  if (status /= NF90_NOERR) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting number of interfaces of "//&
      trim(int_depth_name)//" in "//trim(depth_file))
    nz_out = -1 ; return
  endif

  if (ni < 2) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
        "At least two interface depths must be specified in "//trim(depth_file))
    nz_out = -1 ; return
  endif

  status = NF90_INQ_DIMID(ncid, cell_depth_name, layid)
  if (status /= NF90_NOERR) call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting ID of dimension "//&
      trim(cell_depth_name)//" in "//trim(depth_file))
  status = nf90_Inquire_Dimension(ncid, layid, len=nz_out)
  if (status /= NF90_NOERR) call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting number of interfaces of "//&
      trim(cell_depth_name)//" in "//trim(depth_file))
  if (ni /= nz_out+1) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
        "The interface depths must have one more point than cell centers in "//&
        trim(depth_file))
    nz_out = -1 ; return
  endif

  allocate(int_depth(nz_out+1))
  allocate(cell_depth(nz_out))

  status = NF90_INQ_VARID(ncid, int_depth_name, intvid)
  if (status /= NF90_NOERR) call GOLD_error(FATAL,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting ID of variable "//&
      trim(int_depth_name)//" in "//trim(depth_file))
  status = NF90_GET_VAR(ncid, intvid, int_depth)
  if (status /= NF90_NOERR) call GOLD_error(FATAL,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Reading variable "//&
      trim(int_depth_name)//" in "//trim(depth_file))
  status = NF90_GET_ATT(ncid, intvid, "units", units)
  if (status /= NF90_NOERR) units = "meters"
  status = NF90_GET_ATT(ncid, intvid, "long_name", long_name)
  if (status /= NF90_NOERR) long_name = "Depth of edges"
  edge_index = diag_axis_init(int_depth_name, int_depth, units, 'z', &
                              long_name, direction=-1)

! Create an fms axis with the same data as the cell_depth array in the input file.
  status = NF90_INQ_VARID(ncid, cell_depth_name, layvid)
  if (status /= NF90_NOERR) call GOLD_error(FATAL,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Getting ID of variable "//&
      trim(cell_depth_name)//" in "//trim(depth_file))
  status = NF90_GET_VAR(ncid, layvid, cell_depth)
  if (status /= NF90_NOERR) call GOLD_error(FATAL,"GOLD_diag_to_Z get_Z_depths: "//&
      trim(NF90_STRERROR(status))//" Reading variable "//&
      trim(cell_depth_name)//" in "//trim(depth_file))
  status = NF90_GET_ATT(ncid, layvid, "units", units)
  if (status /= NF90_NOERR) units = "meters"
  status = NF90_GET_ATT(ncid, layvid, "long_name", long_name)
  if (status /= NF90_NOERR) long_name = "Depth of cell center"

  z_axis_index = diag_axis_init(cell_depth_name, cell_depth, units, 'z',&
                                long_name, edges = edge_index, direction=-1)
  
  deallocate(cell_depth)

  status = nf90_close(ncid)

  ! Check the sign convention and change to the GOLD "height" convention.
  if (int_depth(1) < int_depth(2)) then 
    do k=1,nz_out+1 ; int_depth(k) = -1*int_depth(k) ; enddo
  endif

  ! Check for inversions in grid.
  do k=1,nz_out ; if (int_depth(k) < int_depth(k+1)) then
    call GOLD_error(WARNING,"GOLD_diag_to_Z get_Z_depths: "//&
        "Inverted interface depths in output grid in "//depth_file)
    nz_out = -1 ; deallocate(int_depth) ; return
  endif ; enddo

end subroutine get_Z_depths

subroutine GOLD_diag_to_Z_end(CS)
  type(diag_to_Z_CS), pointer :: CS
  integer :: m

  if (ASSOCIATED(CS%u_z))   deallocate(CS%u_z)
  if (ASSOCIATED(CS%v_z))   deallocate(CS%v_z)
  if (ASSOCIATED(CS%Z_int)) deallocate(CS%Z_int)
  do m=1,CS%num_tr_used ;   deallocate(CS%tr_z(m)%p) ; enddo

  deallocate(CS)

end subroutine GOLD_diag_to_Z_end

function ocean_register_diag_with_z (tr_ptr, vardesc_tr, G, Time, CS)
  real, dimension(NXMEM_,NYMEM_,NZ_), target, intent(in) :: tr_ptr
  type(vardesc),                              intent(in) :: vardesc_tr
  type(ocean_grid_type),                      intent(in) :: G
  type(time_type),                            intent(in) :: Time
  type(diag_to_Z_CS),                         pointer    :: CS
  integer                                     :: ocean_register_diag_with_z
!   This subroutine registers a tracer to be output in depth space.
! Arguments: tr_ptr - the tracer that is being offered for translation to Z-space.
!  (in)      vardesc_tr - The descriptor of the variable.
!  (in)      Time - The current model time.
!  (in)      G - The ocean's grid structure.
!  (in)      CS - The control structure returned by a previous call to
!                 diagnostics_init.
  type(vardesc) :: vardesc_z
  integer :: isd, ied, jsd, jed, nz, m, id_test
  isd = G%isd ; ied = G%ied ; jsd = G%jsd ; jed = G%jed ; nz = G%ke
  if (.not.associated(CS)) call GOLD_error(FATAL, &
         "register_Z_tracer: Module must be initialized before it is used.")

  if (CS%num_tr_used >= MAX_FIELDS) then
    call GOLD_error(WARNING,"ocean_register_diag_with_z:  Attempted to register and use "//&
                   "more than MAX_FIELDS z-space tracers via ocean_register_diag_with_z.")
    return
  endif

! Register the layer tracer
  ocean_register_diag_with_z =  ocean_register_diag(vardesc_tr, G, Time)

! Copy the layer tracer variable descriptor to a z-tracer descriptor
! Change the name and layer information.
  vardesc_z = vardesc_tr
  vardesc_z%name = trim(vardesc_tr%name)//"_z"
  vardesc_z%z_grid = "z"
  m = CS%num_tr_used + 1
  CS%missing_tr(m) = CS%missing_value ! This could be changed later, if desired.
  CS%id_tr(m) =  register_Z_diag(vardesc_z, CS, Time, CS%missing_tr(m))

  if (CS%nz_zspace > NO_ZSPACE) then
! There is a depth-space target file.
    if (CS%id_tr(m)>0) then
! Only allocate the tr_z field id there is a diag_table entry looking
! for it.
      CS%num_tr_used = m
      call safe_alloc_ptr(CS%tr_z(m)%p,isd,ied,jsd,jed,CS%nz_zspace)
!Can we do the following at this point?
! tr_ptr might not be allocated yet
      CS%tr_model(m)%p => tr_ptr
    endif
  else
! There is no depth-space target file but warn if a diag_table entry is
! present.
    if (CS%id_tr(m)>0) call GOLD_error(WARNING, &
        "ocean_register_diag_with_z: "//trim(vardesc_z%name)// &
        " cannot be output without an appropriate depth-space target file.")
  endif

end function ocean_register_diag_with_z

function register_Z_diag(var_desc, CS, day, missing)
  integer :: register_Z_diag
  type(vardesc),       intent(in) :: var_desc
  type(diag_to_Z_CS),  pointer    :: CS
  type(time_type),     intent(in) :: day
  real,                intent(in) :: missing

  integer, dimension(3)  :: axes

  ! Use the hor_grid and z_grid components of vardesc to determine the 
  ! desired axes to register the diagnostic field for.
  select case (var_desc%z_grid)

    case ("z")
      select case (var_desc%hor_grid)
        case ("q")
          axes = CS%axesqz
        case ("h")
          axes = CS%axeshz
        case ("u")
          axes = CS%axesuz
        case ("v")
          axes = CS%axesvz
        case default
          call GOLD_error(FATAL,&
            "register_Z_diag: unknown hor_grid component "//trim(var_desc%hor_grid))
      end select

    case default
        call GOLD_error(FATAL,&
          "register_Z_diag: unknown z_grid component "//trim(var_desc%z_grid))
  end select

  register_Z_diag = register_diag_field("ocean_model", trim(var_desc%name), axes, &
        day, trim(var_desc%longname), trim(var_desc%units), missing_value=missing)

end function register_Z_diag

function register_Zint_diag(var_desc, CS, day)
  integer :: register_Zint_diag
  type(vardesc),       intent(in) :: var_desc
  type(diag_to_Z_CS),  pointer    :: CS
  type(time_type),     intent(in) :: day
  integer, dimension(3)  :: axes

  if (CS%nz_zspace < 0) then
    register_Zint_diag = -1 ; return
  endif

  ! Use the hor_grid and z_grid components of vardesc to determine the 
  ! desired axes to register the diagnostic field for.
  select case (var_desc%hor_grid)
    case ("h")
      axes = CS%axeshzi
    case ("q")
      axes = CS%axesqzi
    case ("u")
      axes = CS%axesuzi
    case ("v")
      axes = CS%axesvzi
    case default
      call GOLD_error(FATAL,&
        "register_Z_diag: unknown hor_grid component "//trim(var_desc%hor_grid))
  end select

  register_Zint_diag = register_diag_field("ocean_model", trim(var_desc%name), &
        axes, day, trim(var_desc%longname), trim(var_desc%units), &
        missing_value=CS%missing_value)

end function register_Zint_diag


end module GOLD_diag_to_Z
